#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import torch

from .base_checkpointer import BaseCheckpointerPlugin


class SamplePrintingCheckpointerPlugin(BaseCheckpointerPlugin):
    """样本打印保存点插件 - 打印样本信息但不保存检查点"""
    
    def __init__(self, config):
        """
        初始化样本打印保存点插件
        
        Args:
            config (dict): 配置参数字典
        """
        super().__init__(config)
        # 设置默认的打印间隔
        self.print_interval = config.get('sample_printing_interval', 100)
        self.step_count = 0
        
        # 检查是否是rank 0进程
        self.is_rank_zero = self._is_rank_zero()
        
    def _is_rank_zero(self):
        """检查是否是rank 0进程"""
        try:
            # 首先尝试从环境变量获取
            if "RANK" in os.environ:
                return int(os.environ["RANK"]) == 0
                
            # 默认返回True（非分布式训练）
            return True
        except:
            return True
        
    def should_save_checkpoint(self, current_loss, avg_loss, epoch, step):
        """
        判断是否应该保存检查点
        
        Args:
            current_loss (float): 当前损失值
            avg_loss (float): 平均损失值
            epoch (int): 当前epoch
            step (int): 当前step
            
        Returns:
            bool: 是否应该保存检查点
        """
        self.step_count += 1
        
        # 每隔一定步数打印样本信息，且仅在rank 0进程中执行
        if self.is_rank_zero and self.step_count % self.print_interval == 0:
            print(f"[SAMPLE] Epoch: {epoch}, Step: {step}, Current Loss: {current_loss:.4f}, Avg Loss: {avg_loss:.4f}")
            # 与noop插件保持一致，始终返回False，不触发保存检查点操作
            # 样本打印功能通过print_sample方法在训练器中直接调用
            return False
            
        return False
        
    def should_early_stop(self):
        """
        判断是否应该早停
        
        Returns:
            bool: 是否应该早停
        """
        # 此插件不处理早停
        return False
        
    def print_sample(self, model, tokenizer, device):
        """
        打印生成的样本
        
        Args:
            model: 模型实例
            tokenizer: 分词器实例
            device: 设备
        """
        # 仅在rank 0进程中执行
        if not self.is_rank_zero:
            return
            
        try:
            # 准备输入文本
            input_text = "人工智能是"
            inputs = tokenizer(input_text, return_tensors='pt')
            input_ids = inputs['input_ids']
            
            # 为避免设备的兼容性问题，将样本生成移到CPU上执行
            input_ids = input_ids.to('cpu')
            
            # 创建正确类型的attention_mask
            attention_mask = torch.ones_like(input_ids, dtype=torch.float32)
                
            # 临时创建一个模型的副本来进行推理
            # 处理分布式模型的情况
            if hasattr(model, 'module'):
                # 如果是分布式模型（DistributedDataParallel包装的模型）
                model_to_copy = model.module
            else:
                model_to_copy = model
            
            # 确保在创建副本时不干扰原始模型的状态
            with torch.no_grad():
                model_copy = type(model_to_copy)(model_to_copy.config).cpu()
                model_copy.load_state_dict(model_to_copy.state_dict())
                    
                # 生成文本
                output = model_copy.generate(
                    input_ids,
                    attention_mask=attention_mask,
                    max_length=50,
                    num_return_sequences=1,
                    do_sample=True,
                    temperature=0.7,
                    pad_token_id=tokenizer.eos_token_id
                )
                
            # 解码生成的文本
            generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
            print(f"[SAMPLE GENERATION] 输入: {input_text}")
            print(f"[SAMPLE GENERATION] 输出: {generated_text}")
            
            # 清理副本
            del model_copy
            
        except Exception as e:
            print(f"[WARNING] 生成样本时出错: {e}")